import pandas as pd
import numpy as np
import json
import sys
import warnings
warnings.filterwarnings('ignore')

def select_top_features_with_mandatory(correlations, mandatory_fields=['rcyl', 'rcql', 'ljcyl', 'ljcql', 'qyb'], max_features=20):
    """
    选择前20个特征，确保包含必须字段，并且这5个字段尽量排在前面
    """
    mandatory_fields_lower = [field.lower() for field in mandatory_fields]

    # 按相关性排序所有字段
    sorted_features = sorted(correlations.items(), key=lambda x: abs(x[1]), reverse=True)

    selected = []
    mandatory_correlations = []
    other_correlations = []

    # 分离必须字段和其他字段
    for field, corr in sorted_features:
        if field.lower() in mandatory_fields_lower:
            mandatory_correlations.append((field, corr))
        else:
            other_correlations.append((field, corr))

    # 必须字段按相关性排序
    mandatory_correlations.sort(key=lambda x: abs(x[1]), reverse=True)

    # 先添加必须字段
    mandatory_added = []
    for field, corr in mandatory_correlations:
        selected.append(field)
        mandatory_added.append(field)

    # 然后添加其他高相关性字段，确保总数不超过max_features
    for field, corr in other_correlations:
        if len(selected) >= max_features:
            break
        selected.append(field)

    return selected, mandatory_added

def main(csv_file_path):
    try:
        print(f"开始第一次相关性分析，数据文件: {csv_file_path}")

        # 读取CSV数据
        df = pd.read_csv(csv_file_path)
        print(f"数据形状: {df.shape}")

        # 确保qyb字段存在
        if 'qyb' not in df.columns:
            result = {
                'status': 'error',
                'message': 'qyb字段不存在于数据中'
            }
            print(json.dumps(result))
            return

        # 获取所有数值字段
        numeric_features = df.select_dtypes(include=[np.number]).columns.tolist()

        print(f"找到 {len(numeric_features)} 个数值特征字段")

        # 计算与qyb的相关性（包括qyb自身，相关性为1.0）
        correlations_with_qyb = {}
        qyb_series = df['qyb'].dropna()

        for feature in numeric_features:
            if feature == 'qyb':
                # qyb与自身的相关性为1.0
                correlations_with_qyb[feature] = 1.0
                print(f"{feature}: 相关性 = 1.0000 (目标变量)")
                continue

            feature_series = df[feature].dropna()

            # 确保有足够的非空值进行相关性计算
            if len(feature_series) > 1 and len(qyb_series) > 1:
                # 找到两个序列的交集索引
                common_indices = qyb_series.index.intersection(feature_series.index)
                if len(common_indices) > 1:
                    corr = qyb_series.loc[common_indices].corr(feature_series.loc[common_indices])
                    if not np.isnan(corr):
                        correlations_with_qyb[feature] = abs(corr)
                        print(f"{feature}: 相关性 = {abs(corr):.4f}")

        print(f"成功计算了 {len(correlations_with_qyb)} 个字段的相关性")

        # 定义必须包含的字段
        mandatory_fields = ['rcyl', 'rcql', 'ljcyl', 'ljcql', 'qyb']

        # 选择前20个特征，确保包含必须字段
        top_20_features, mandatory_included = select_top_features_with_mandatory(
            correlations_with_qyb, mandatory_fields, 20)

        print(f"选择的前20个特征: {top_20_features}")
        print(f"必须字段包含情况: {mandatory_included}")

        # 检查缺失的必须字段
        missing_mandatory = []
        for field in mandatory_fields:
            field_found = False
            for included_field in mandatory_included:
                if included_field.lower() == field.lower():
                    field_found = True
                    break
            if not field_found:
                missing_mandatory.append(field)

        # 返回结果
        result = {
            'top_20_features': top_20_features,
            'correlations': correlations_with_qyb,
            'mandatory_fields': mandatory_fields,
            'mandatory_included': mandatory_included,
            'missing_mandatory': missing_mandatory,
            'total_features_analyzed': len(numeric_features),
            'valid_correlations_count': len(correlations_with_qyb),
            'data_shape': df.shape,
            'qyb_non_null_count': len(qyb_series),
            'target_variable': 'qyb',
            'status': 'success'
        }

        print(json.dumps(result))

    except Exception as e:
        error_result = {
            'status': 'error',
            'message': f'第一次相关性分析失败: {str(e)}'
        }
        print(json.dumps(error_result))

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(json.dumps({
            'status': 'error',
            'message': '参数错误，需要提供CSV文件路径'
        }))
    else:
        main(sys.argv[1])
